import numpy as np

from utils.bbox_utils import xywh2xyxy
from utils.IoU_utils import get_IoU

# 暂存区，之后写到配置文件里
INPUT_SIZE = (640, 640)


def resizedImage2YoloInput(resized_image: np.ndarray):
    img = np.transpose(resized_image, (2, 0, 1))  # h, w, c -> c, h, w
    input_tensor = img / 255
    return np.expand_dims(input_tensor, axis=0).astype(np.float32)  # c, h, w -> b, c, h, w


def output_process(output_tensor: np.ndarray):
    # (1, 4 + class_num, 8400) -> (4 + class_num, 8400)
    output_tensor = np.squeeze(output_tensor, axis=0)
    # (4 + class_num, 8400) -> (8400, 4 + class_num)
    output_tensor = np.transpose(output_tensor, (1, 0))
    predict_class_list = output_tensor[..., 4:]
    predict_conf = np.max(predict_class_list, axis=-1)
    # (8400, 4 + class_num) -> (8400, 4 + 1 + class_num)
    output_tensor = np.insert(output_tensor, 4, predict_conf, axis=-1)
    return output_tensor  # (8400, 4 + 1 + class_num)


def confidence_filter(output_tensor: np.ndarray, conf_threshold):
    """置信度过滤"""
    return output_tensor[output_tensor[..., 4] > conf_threshold]


def NMS(output_tensor: np.ndarray, IoU_threshold):
    predict_class_conf = output_tensor[..., 5:]
    cls = []
    for i in range(len(predict_class_conf)):
        cls.append(int(np.argmax(predict_class_conf[i])))

    total_cls = list(set(cls))
    output_box = []

    for i in range(len(total_cls)):
        current_cls = total_cls[i]
        current_cls_box = []
        temp = output_tensor[:, :6]
        for j in range(len(cls)):
            if cls[j] == current_cls:
                temp[j][5] = current_cls
                current_cls_box.append(temp[j])
        current_cls_box = np.array(current_cls_box)
        sorted_cls_box = sorted(current_cls_box, key=lambda x: -x[4])
        max_conf_box = sorted_cls_box[0]
        output_box.append(max_conf_box)
        sorted_cls_box = np.delete(sorted_cls_box, 0, 0)
        while len(sorted_cls_box) > 0:
            max_conf_box = output_box[-1]
            del_index = []
            for k in range(len(sorted_cls_box)):
                current_box = sorted_cls_box[k]
                IoU = get_IoU(max_conf_box, current_box)
                if IoU > IoU_threshold:
                    del_index.append(k)
            sorted_cls_box = np.delete(sorted_cls_box, del_index, 0)
            if len(sorted_cls_box) > 0:
                output_box.append(sorted_cls_box[0])
                sorted_cls_box = np.delete(sorted_cls_box, 0, 0)
    return np.array(output_box)


def re_letterbox(output: np.ndarray, img: np.ndarray):
    if len(output) == 0:
        return None
    x, y, w, h, conf, cls = output.transpose((1, 0))
    x1, y1, x2, y2 = xywh2xyxy(x, y, w, h)

    before_height, before_width, _ = img.shape
    after_height, after_width = INPUT_SIZE

    scale = max(before_height / after_height, before_width / after_width)

    new_height, new_width = before_height / scale, before_width / scale
    x_move, y_move = abs(new_width - after_width) // 2, abs(new_height - after_height) // 2
    n_x1, n_x2 = (x1 - x_move) * scale, (x2 - x_move) * scale
    n_y1, n_y2 = (y1 - y_move) * scale, (y2 - y_move) * scale
    n_xyxy = np.array([n_x1, n_y1, n_x2, n_y2, conf, cls]).transpose((1, 0))
    return n_xyxy
